# models/prototype.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from models.enhanced_cnn_transformer import EnhancedCNNTransformer


class PrototypeEnhancedCNNTransformer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.base_model = EnhancedCNNTransformer(config)

        # 修改分类头以输出原型特征
        in_features = 64 * 3 * 30  # 来自CNN-Transformer的输出维度
        self.prototype_layer = nn.Linear(in_features, config.prototype_dim)

        # 原型分类头
        self.classifier = nn.Sequential(
            nn.Linear(config.prototype_dim, 512),
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(512, config.num_classes)
        )

        # 可学习的类别原型
        self.register_buffer('class_prototypes',
                             torch.zeros(config.num_classes, config.prototype_dim))

    def forward(self, x, labels=None):
        # CNN特征提取 [batch, C, 3, 30, T]
        features = self.base_model.cnn(x)

        # 重组为时序数据
        batch_size = features.size(0)
        features = features.permute(0, 4, 1, 2, 3)  # [batch, T, C, 3, 30]
        features = features.reshape(batch_size, -1, 64 * 3 * 30)  # [batch, T, D]

        # Transformer处理
        features = self.base_model.transformer_encoder(features)
        features = features.mean(dim=1)  # 全局平均

        # 原型特征
        prototype_features = self.prototype_layer(features)
        logits = self.classifier(prototype_features)

        # 计算原型损失
        prototype_loss = 0.0
        if labels is not None and self.training:
            # 更新类别原型 (指数移动平均)
            with torch.no_grad():
                for class_idx in range(self.config.num_classes):
                    mask = (labels == class_idx)
                    if mask.any():
                        class_features = prototype_features[mask].mean(dim=0)
                        self.class_prototypes[class_idx] = (
                                0.9 * self.class_prototypes[class_idx] +
                                0.1 * class_features
                        )

            # 计算对比损失
            prototype_loss = self.prototype_contrastive_loss(
                prototype_features, labels
            )

        return logits, prototype_loss, prototype_features

    def prototype_contrastive_loss(self, features, labels):
        # 归一化特征和原型
        features = F.normalize(features, dim=1)
        prototypes = F.normalize(self.class_prototypes, dim=1)

        # 计算相似度矩阵
        logits = torch.mm(features, prototypes.t()) / self.config.temperature

        # 对比损失
        loss = F.cross_entropy(logits, labels)
        return loss